use std::collections::HashMap;
use std::fs;
use serde_json::Value;

struct UnionFind {
    parent: HashMap<String, String>,
    rank: HashMap<String, i32>,
}

impl UnionFind {
    fn new() -> Self {
        UnionFind {
            parent: HashMap::new(),
            rank: HashMap::new(),
        }
    }

    fn find(&mut self, x: &str) -> String {
        if !self.parent.contains_key(x) {
            self.parent.insert(x.to_string(), x.to_string());
            self.rank.insert(x.to_string(), 0);
            return x.to_string();
        }
        if self.parent[x] != x {
            let parent = self.parent[x].clone();
            let root = self.find(&parent);
            self.parent.insert(x.to_string(), root.clone());
            root
        } else {
            x.to_string()
        }
    }

    fn union(&mut self, x: &str, y: &str) {
        let x_root = self.find(x);
        let y_root = self.find(y);
        if x_root == y_root {
            return;
        }
        let x_rank = self.rank[&x_root];
        let y_rank = self.rank[&y_root];
        if x_rank < y_rank {
            self.parent.insert(x_root, y_root);
        } else if x_rank > y_rank {
            self.parent.insert(y_root, x_root);
        } else {
            self.parent.insert(y_root.clone(), x_root.clone());
            self.rank.insert(x_root, x_rank + 1);
        }
    }

    fn add(&mut self, x: &str) {
        self.find(x);
    }
}

pub fn count_provinces() -> String {
    let data = fs::read_to_string("district.json").expect("Unable to read file");
    let json: Value = serde_json::from_str(&data).expect("Unable to parse json");
    let mut ans: Vec<usize> = Vec::new();
    let mut uf = UnionFind::new();

    for batch in json.as_object().unwrap().values() {
        let mut city_neighbors: HashMap<String, Vec<String>> = HashMap::new();
        for (city, neighbors) in batch.as_object().unwrap() {
            if *city==neighbors[0] {
                continue;
            }
            city_neighbors.entry(city.clone()).or_insert_with(Vec::new).extend(
                neighbors.as_array().unwrap().iter().map(|n| n.as_str().unwrap().to_string())
            );
        }
        for (city, neighbors) in city_neighbors.iter() {
            uf.add(city);
            for neighbor in neighbors {
                uf.add(neighbor);
                uf.union(city, neighbor);
            }
        }
        let mut provinces = std::collections::HashSet::new();
        let keys: Vec<String> = uf.parent.keys().cloned().collect();
        for city in keys {
            provinces.insert(uf.find(&city));
        }
        ans.push(provinces.len());
        // 清空并初始化并查集
        uf = UnionFind::new();
    }
    ans.iter().map(|x| x.to_string()).collect::<Vec<String>>().join(",")
}

